import Core.IO as IO
import Core.Gadget as Gadget
import pymongo
import pandas as pd

#---Line1: ---
#---Line2: ---
def PreparePlotData(database, name, datetime1, datetime2, instrumentType = "Stock"):
    #
    filter = {}
    filter["StdDateTime"] = {"$gte": datetime1, "$lte": datetime2}
    sort = [("StdDateTime", pymongo.ASCENDING)]
    #
    if instrumentType == "Portfolio":
        filter["Portfolio"] = name
        equityPrices = database.findWithFilter("Portfolio", "Account", filter, sort)

        # ---To DataFrame---
        dfEquityPrices = IO.DataListToDataFrame(equityPrices, ["StdDateTime", "UnitNetValue"])
        dfEquityPrices.rename(columns={'UnitNetValue': "Close"}, inplace=True)

    elif instrumentType == "Stock":
        equityPrices = database.findWithFilter("Quote", name + "_Time_86400_Bar", filter, sort)

        # ---To DataFrame---
        dfEquityPrices = IO.DataListToDataFrame(equityPrices, ["StdDateTime", "BClose"])
        dfEquityPrices.rename(columns={'BClose': "Close"}, inplace=True)
    #
    if len(equityPrices) == 0:
        return None

    # --- Refresh Beginning Ending DateTimes ---
    # --- Equity / Portfolio 的起始时间可能小于Benchmark ---
    datetime1 = equityPrices[0]["StdDateTime"]
    datetime2 = equityPrices[len(equityPrices) - 1]["StdDateTime"]

    # ---Benchmark---
    bmSymbol = "000001.SH"
    bmPrices = database.find("Index", bmSymbol + "_Time_86400_Bar", beginDateTime=datetime1, endDateTime=datetime2)
    dfBenchmarkPrices = IO.DataListToDataFrame(bmPrices, ["StdDateTime", "Close"])
    dfBenchmarkPrices.rename(columns={'Close': bmSymbol}, inplace=True)

    # --- Merge & 补齐---
    dfCombine = pd.merge(dfBenchmarkPrices, dfEquityPrices, on='StdDateTime', how='outer')
    dfCombine.fillna(method='ffill', inplace=True)

    # ---Calculate---
    dfCombine["BM"] = dfCombine[bmSymbol] / dfCombine.loc[0,bmSymbol]
    dfCombine["Equity"] = dfCombine["Close"] / dfCombine.loc[0, "Close"]

    #
    prices = {}
    prices["X"] = []
    prices["Y1"] = []
    prices["Y2"] = []
    bmBase = 1
    equityBase = 1
    #

    #
    preDateTime = None
    for i in range(len(dfCombine)):
        curDateTime = dfCombine.loc[i, "StdDateTime"]
        # record Beginning and Ending
        insert = False
        if i == 0 or i == len(dfCombine)-1:
            insert = True

        # monthly update
        if preDateTime != None  and preDateTime.month != curDateTime.month:
            insert = True

        if insert:
            sDate = Gadget.ToDateString(curDateTime)
            prices["X"].append(sDate)
            prices["Y1"].append(dfCombine.loc[i, "BM"])
            prices["Y2"].append(dfCombine.loc[i, "Equity"])

        preDateTime = curDateTime
    #
    return prices